from rank_bm25 import BM25Okapi
from retrieval.reranker import CrossEncoderReranker, RRFReranker


class HybridRetriever:
    def __init__(self, vector_db):
        self.vector_db = vector_db
        self.bm25_index = self._build_bm25_index()
        self.model_reranker = CrossEncoderReranker()
        self.RRF_reranker = RRFReranker()

    def _build_bm25_index(self):
        """从向量数据库构建BM25索引"""
        # 获取所有文档
        documents = [doc for doc in self.vector_db.get()["documents"] if doc.strip()]

        # 分词处理
        tokenized_docs = [doc.split() for doc in documents]
        return BM25Okapi(tokenized_docs)

    def model_retrieve(self, query, top_n=5, rerank_top_k=20):
        """执行混合检索（语义+关键词）"""
        # 语义检索 (Vector)
        vector_results = self.vector_db.similarity_search(query, k=rerank_top_k)
        vector_docs = [doc.page_content for doc in vector_results]

        # 关键词检索 (BM25)
        tokenized_query = query.split()
        bm25_scores = self.bm25_index.get_scores(tokenized_query)
        top_indices = bm25_scores.argsort()[-rerank_top_k:][::-1]
        bm25_docs = [self.vector_db.get()["documents"][i] for i in top_indices]

        # 融合结果并去重
        combined_docs = list(set(vector_docs + bm25_docs))

        # 重排序提升相关性
        reranked_docs = self.model_reranker.rerank(query, combined_docs, top_n)

        return reranked_docs

    def RRF_retrieve(self, query, top_n=5, rerank_top_k=20):
        """执行混合检索（语义+关键词）并使用RRF重排序"""
        # 语义检索 (Vector)
        vector_results = self.vector_db.similarity_search(query, k=rerank_top_k)
        vector_docs = [doc.page_content for doc in vector_results]

        # 关键词检索 (BM25)
        tokenized_query = query.split()
        bm25_scores = self.bm25_index.get_scores(tokenized_query)
        top_indices = bm25_scores.argsort()[-rerank_top_k:][::-1]
        bm25_docs = [self.vector_db.get()["documents"][i] for i in top_indices]

        # 使用RRF重排序
        reranked_docs = self.RRF_reranker.rerank(vector_docs, bm25_docs, top_n)

        return reranked_docs